import json
import pandas as pd
import numpy as np
from typing import Dict, List, Optional
import re


VALID_SCHOOLS = {
    '985': [
        '清华大学', '北京大学', '浙江大学', '上海交通大学', '复旦大学',
        '南京大学', '中国科学技术大学', '哈尔滨工业大学', '西安交通大学',
        '华中科技大学', '武汉大学', '吉林大学', '四川大学', '山东大学',
        '同济大学', '华南理工大学', '中山大学', '北京理工大学',
        '天津大学', '东南大学', '南开大学', '厦门大学', '大连理工大学',
        '华东师范大学', '中南大学', '湖南大学', '北京航空航天大学',
        '重庆大学', '兰州大学', '电子科技大学', '东北大学',
        '中国人民大学', '北京师范大学', '中国农业大学', '中央民族大学',
        '西北工业大学', '西北农林科技大学', '中国海洋大学'
    ],
    '211': [
        '清华大学', '北京大学', '浙江大学', '上海交通大学', '复旦大学',
        '南京大学', '中国科学技术大学', '哈尔滨工业大学', '西安交通大学',
        '华中科技大学', '武汉大学', '吉林大学', '四川大学', '山东大学',
        '同济大学', '华南理工大学', '中山大学', '北京理工大学',
        '天津大学', '东南大学', '南开大学', '厦门大学', '大连理工大学',
        '华东师范大学', '中南大学', '湖南大学', '北京航空航天大学',
        '重庆大学', '兰州大学', '电子科技大学', '东北大学',
        '中国人民大学', '北京师范大学', '中国农业大学', '中央民族大学',
        '西北工业大学', '西北农林科技大学', '中国海洋大学',
        '北京邮电大学', '北京工业大学', '北京科技大学', '北京化工大学',
        '北京交通大学', '华北电力大学', '北京林业大学', '北京中医药大学',
        '北京协和医学院', '北京外国语大学', '中国政法大学', '首都医科大学',
        '南京理工大学', '南京航空航天大学', '苏州大学', '南京农业大学',
        '中国药科大学', '华东理工大学', '东华大学', '上海大学',
        '上海财经大学', '上海外国语大学', '暨南大学', '华南师范大学',
        '郑州大学', '西北大学', '西南交通大学', '广西大学', '云南大学',
        '贵州大学', '新疆大学', '石河子大学', '西藏大学', '宁夏大学',
        '青海大学', '内蒙古大学', '海南大学', '河北工业大学', '太原理工大学',
        '延边大学', '辽宁大学', '大连海事大学', '东北师范大学', '哈尔滨工程大学',
        '东北林业大学', '东北农业大学', '福州大学', '合肥工业大学', '安徽大学',
        '南昌大学', '河南大学', '武汉理工大学', '华中师范大学', '华中农业大学',
        '中南财经政法大学', '湖南师范大学', '国防科技大学', '四川农业大学',
        '西南财经大学', '重庆大学', '西南大学', '贵州大学', '云南大学',
        '西北大学', '西安电子科技大学', '陕西师范大学', '空军军医大学', '长安大学',
        '中国矿业大学', '中国地质大学', '河海大学', '江南大学', '南京师范大学',
        '中国石油大学'
    ]
}

def normalize_school(school: str) -> str:
    """标准化学校名称"""
    if not school or school == '无':
        return '普通'

    school_clean = school.replace('大学', '').replace('学院', '')

    for level, schools in VALID_SCHOOLS.items():
        for valid_school in schools:
            # 检查是否匹配
            valid_school_clean = valid_school.replace('大学', '').replace('学院', '')
            if school_clean in valid_school_clean or valid_school_clean in school_clean:
                return level
                
    return '普通'

def extract_level_value(level: str) -> int:
    """从职级字符串中提取数值"""
    if not level:
        return 0
    # 使用正则表达式提取数字
    match = re.search(r'\d+', level)
    return int(match.group()) if match else 0

def extract_tech_info(focuses: List[str], tech_stacks: List[str]) -> Dict:
    """从focuses和tech_stacks中提取技术相关信息"""
    # 技术方向映射
    direction_keywords = {
        '后端': ['后端', '服务端', 'backend', 'server'],
        '前端': ['前端', 'frontend', 'web'],
        '算法': ['算法', '机器学习', '深度学习', 'AI'],
        '数据': ['数据', '大数据', 'data'],
        '测试': ['测试', 'test', 'QA'],
        '运维': ['运维', 'devops', 'SRE'],
    }
    
    # 技术领域映射
    domain_keywords = {
        '云计算': ['云', '云计算', 'cloud'],
        'AI': ['AI', '机器学习', '深度学习', '算法'],
        '大数据': ['大数据', '数据分析', 'hadoop', 'spark'],
        '基础架构': ['架构', '中间件', '基础设施'],
        '应用开发': ['应用', 'app', '开发', 'web'],
    }
    
    # 默认值
    tech_info = {
        'tech_direction': 'unknown',
        'tech_domain': 'unknown',
        'main_tech': 'unknown'
    }
    
    # 合并focuses和tech_stacks
    all_tech_info = ' '.join((focuses or []) + (tech_stacks or [])).lower()
    
    # 判断技术方向
    for direction, keywords in direction_keywords.items():
        if any(keyword.lower() in all_tech_info for keyword in keywords):
            tech_info['tech_direction'] = direction
            break
            
    # 判断技术领域
    for domain, keywords in domain_keywords.items():
        if any(keyword.lower() in all_tech_info for keyword in keywords):
            tech_info['tech_domain'] = domain
            break
            
    # 提取主技术栈
    if tech_stacks and len(tech_stacks) > 0:
        tech_info['main_tech'] = tech_stacks[0]
    
    return tech_info

def determine_school_level(bachelor_school: str, advanced_school: str) -> str:
    """判断学校等级"""

    bachelor_level = normalize_school(bachelor_school)
    advanced_level = normalize_school(advanced_school)
    
    # 研究生学校等级优先级更高
    if advanced_level in ['985', '211']:
        return advanced_level
    elif bachelor_level in ['985', '211']:
        return bachelor_level
    else:
        # 如果都不是985/211，返回普通
        return '普通'

def calculate_experience_score(year_of_experience: float, year_in_company: float) -> float:
    """计算经验评分"""
    # 总经验权重0.7，公司经验权重0.3
    # P8一般要求8年+经验，所以将经验上限设为10年
    total_exp_score = min(year_of_experience * 7, 70)  # 最高70分
    company_exp_score = min(year_in_company * 6, 30)   # 最高30分
    return total_exp_score + company_exp_score

def calculate_education_score(degree: str, school_level: str) -> float:
    """计算教育评分"""
    # 学历分数 (最高60分)
    degree_scores = {
        '博士': 60,
        '硕士': 50,
        '本科': 40,
        '大专': 30,
        'unknown': 20
    }
    
    # 学校等级分数 (最高40分)
    school_scores = {
        '985': 40,
        '211': 35,
        '双非': 30,
        '普通': 25,
        'unknown': 20
    }
    
    degree_score = degree_scores.get(degree, 20)
    school_score = school_scores.get(school_level, 20)
    
    # P8的例子显示本科也可以达到高级别，所以降低学历权重
    weighted_degree_score = degree_score * 0.7  # 降低学历权重
    weighted_school_score = school_score * 0.7  # 降低学校权重
    
    return weighted_degree_score + weighted_school_score

def calculate_technical_score(tech_stacks: Optional[List[str]], focuses: Optional[List[str]]) -> float:
    """计算技术评分"""
    tech_stacks = tech_stacks or []
    focuses = focuses or []
    
    # 技术栈广度分数 (最高30分)
    breadth_score = min(len(tech_stacks) * 10, 30)
    
    # 技术深度分数 (最高40分)
    depth_keywords = ['算法', '机器学习', '人工智能', '架构', '中间件', '基础设施', '分布式', '云原生']
    depth_score = 0
    for focus in focuses:
        if any(keyword in focus.lower() for keyword in depth_keywords):
            depth_score += 20
    depth_score = min(depth_score, 40)
    
    # 技术匹配度分数 (最高30分)
    core_tech = ['java', 'python', 'c++', 'go', 'rust', 'typescript']
    match_score = 0
    for tech in tech_stacks:
        if tech.lower() in core_tech:
            match_score += 15
    match_score = min(match_score, 30)
    
    return breadth_score + depth_score + match_score

def calculate_workload_score(hour_per_week: float, work_schedule: str) -> float:
    """计算工作量评分"""
    # 工作时长分数 (最高60分)
    if hour_per_week <= 0:  # 处理缺失值
        if work_schedule == '996':
            hour_per_week = 72
        elif work_schedule == '995':
            hour_per_week = 65
        else:
            hour_per_week = 40
    
    # 标准工作时间是40小时/周
    # 但要考虑加班的实际情况
    if hour_per_week <= 40:
        hours_score = 40
    elif hour_per_week <= 50:
        hours_score = 45
    elif hour_per_week <= 60:
        hours_score = 50
    else:
        hours_score = 60
    
    # 工作制度分数 (最高40分)
    schedule_scores = {
        '996': 40,
        '995': 35,
        '955': 30,
        'unknown': 25
    }
    
    schedule_score = schedule_scores.get(work_schedule, 25)
    
    return hours_score + schedule_score

def calculate_compensation_score(total_comp: float, location: str, base_monthly_comp: float,
                              bonus_comp: float = 0, stock_comp: float = 0) -> float:
    """计算薪酬评分"""
    # 不同城市的P7基准薪酬（以年薪计）
    city_benchmarks = {
        '北京': 800000,
        '上海': 800000,
        '杭州': 750000,
        '深圳': 800000,
        '广州': 750000
    }
    
    benchmark = city_benchmarks.get(location, 700000)
    
    # 总薪酬分数 (最高40分)
    comp_score = min(total_comp / benchmark * 40, 40)
    
    # 月薪分数 (最高20分)
    monthly_benchmark = benchmark / 12
    monthly_score = min(base_monthly_comp / monthly_benchmark * 20, 20)
    
    # 奖金分数 (最高20分)
    bonus_benchmark = benchmark * 0.2  # 假设基准奖金是年薪的20%
    bonus_score = min(bonus_comp / bonus_benchmark * 20, 20)
    
    # 股票分数 (最高20分)
    stock_benchmark = benchmark * 0.3  # 假设基准股票是年薪的30%
    stock_score = min(stock_comp / stock_benchmark * 20, 20)
    
    return comp_score + monthly_score + bonus_score + stock_score

def calculate_responsibility_score(title: str, department: str, total_comp: float) -> float:
    """计算责任评分"""
    # 职位权重分数 (最高40分)
    title = title or ''
    title_score = 20  # 默认分数
    
    title_keywords = {
        '专家': 40,
        '高级': 35,
        '资深': 35,
        '架构': 40,
        '主管': 35,
        '经理': 35
    }
    
    for keyword, score in title_keywords.items():
        if keyword in title:
            title_score = max(title_score, score)
            break
    
    # 部门重要性分数 (最高30分)
    department = department or ''
    dept_score = 15  # 默认分数
    
    core_keywords = ['核心', '中台', '平台', '基础架构', '研发', '算法', 'AI']
    business_keywords = ['业务', '产品', '运营']
    
    if any(keyword in department for keyword in core_keywords):
        dept_score = 30
    elif any(keyword in department for keyword in business_keywords):
        dept_score = 25
    
    # 薪酬等级分数 (最高30分)
    # 使用薪酬水平来推断责任级别
    comp_level_score = 0
    if total_comp >= 1500000:  # P8+
        comp_level_score = 30
    elif total_comp >= 800000:  # P7
        comp_level_score = 25
    elif total_comp >= 500000:  # P6
        comp_level_score = 20
    elif total_comp >= 300000:  # P5
        comp_level_score = 15
    else:
        comp_level_score = 10
    
    return title_score + dept_score + comp_level_score

def determine_location_tier(location: str) -> str:
    """确定城市等级"""
    location = location or 'unknown'  # 处理None值
    tier_1 = ['北京', '上海', '广州', '深圳']
    new_tier_1 = ['成都', '杭州', '重庆', '苏州', '武汉', '西安', '南京', '长沙', '天津', '郑州', '东莞', '佛山', '宁波', '青岛', '合肥']
    
    if location in tier_1:
        return '一线城市'
    elif location in new_tier_1:
        return '新一线'
    else:
        return '其他城市'

def determine_tech_level(title: str, year_of_experience: float) -> str:
    """确定技术等级"""
    title = title or ''  # 处理None值
    if '专家' in title or year_of_experience > 8:
        return '专家'
    elif '高级' in title or year_of_experience > 5:
        return '高级'
    elif '中级' in title or year_of_experience > 3:
        return '中级'
    else:
        return '初级'

def determine_role_type(title: str, team_size: int = 0) -> str:
    """确定角色类型"""
    title = title or ''  # 处理None值
    if team_size > 5 or any(key in title for key in ['主管', '经理', '负责人']):
        return '管理'
    elif any(key in title for key in ['专家', '架构']):
        return '专家'
    else:
        return 'IC'

def determine_business_domain(department: str, focuses: List[str]) -> str:
    """确定业务领域"""
    department = department or ''  # 处理None值
    focuses = focuses or []  # 处理None值
    
    if not department and not focuses:
        return 'unknown'
        
    all_info = department.lower() + ' ' + ' '.join(focuses).lower()
    
    domain_keywords = {
        '电商': ['电商', '商业', 'commerce', '淘宝', '天猫'],
        '金融': ['金融', '支付', 'finance', 'payment'],
        '云计算': ['云', 'cloud', '基础架构'],
        '人工智能': ['ai', '算法', '机器学习', '深度学习'],
        '数据': ['数据', 'data', '分析'],
        '安全': ['安全', 'security', '风控']
    }
    
    for domain, keywords in domain_keywords.items():
        if any(keyword in all_info for keyword in keywords):
            return domain
            
    return '其他'

def standardize_level(level: str) -> str:
    """标准化职级格式"""
    if not level:
        return 'unknown'
        
    # 统一大写
    level = level.upper()
    
    # 处理特殊情况
    if level == 'BAND7':
        return 'P7'
    
    # 确保格式为 P+数字
    if level.startswith('P') and len(level) > 1 and level[1:].isdigit():
        level_num = int(level[1:])
        # 只保留 P4-P9 的数据
        if 4 <= level_num <= 9:
            return level
    
    return 'unknown'

def prepare_ali_data(input_file: str, output_file: str):
    """准备阿里的训练数据"""
    # 读取原始数据
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 筛选阿里的数据
    ali_data = []
    for item in data:
        if item['ugc_data']['company_name'] == '阿里':
            processed_item = {}
            ugc_data = item['ugc_data']
            
            # 标准化职级
            level = standardize_level(ugc_data.get('level', ''))
            if level == 'unknown':
                continue  # 跳过无效职级
                
            processed_item['target_level'] = level
            
            # 计算综合评分特征
            processed_item['experience_score'] = calculate_experience_score(
                float(ugc_data.get('year_of_experience', 0) or 0),
                float(ugc_data.get('year_in_company', 0) or 0)
            )
            
            processed_item['education_score'] = calculate_education_score(
                ugc_data.get('degree', 'unknown'),
                determine_school_level(
                    ugc_data.get('bachelor_school', ''),
                    ugc_data.get('advanced_school', '')
                )
            )
            
            processed_item['technical_score'] = calculate_technical_score(
                ugc_data.get('tech_stacks', []),
                ugc_data.get('focuses', [])
            )
            
            processed_item['workload_score'] = calculate_workload_score(
                float(ugc_data.get('hour_per_week', 0) or 0),
                ugc_data.get('work_schedule', 'unknown')
            )
            
            processed_item['compensation_score'] = calculate_compensation_score(
                float(ugc_data.get('total_comp', 0) or 0),
                ugc_data.get('location', 'unknown'),
                float(ugc_data.get('base_monthly_comp', 0) or 0),
                float(ugc_data.get('bonus_comp', 0) or 0),
                float(ugc_data.get('stock_comp', 0) or 0)
            )
            
            processed_item['responsibility_score'] = calculate_responsibility_score(
                ugc_data.get('title', ''),
                ugc_data.get('department', ''),
                float(ugc_data.get('total_comp', 0) or 0)
            )
            
            # 分类特征
            processed_item['location_tier'] = determine_location_tier(
                ugc_data.get('location', 'unknown')
            )
            
            processed_item['tech_level'] = determine_tech_level(
                ugc_data.get('title', ''),
                float(ugc_data.get('year_of_experience', 0) or 0)
            )
            
            processed_item['role_type'] = determine_role_type(
                ugc_data.get('title', ''),
                0  # 默认team_size为0，因为数据中没有这个字段
            )
            
            processed_item['business_domain'] = determine_business_domain(
                ugc_data.get('department', ''),
                ugc_data.get('focuses', [])
            )
            
            ali_data.append(processed_item)
    
    # 保存处理后的数据
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(ali_data, f, ensure_ascii=False, indent=2)
    
    print(f"共处理了 {len(ali_data)} 条阿里数据")
    
    # 数据分析
    df = pd.DataFrame(ali_data)
    print("\n数据统计信息:")
    print(df.describe())
    print("\n职级分布:")
    print(df['target_level'].value_counts())
    print("\n各评分的相关性:")
    score_columns = [col for col in df.columns if col.endswith('_score')]
    if len(score_columns) > 1:
        correlation_matrix = df[score_columns].corr(method='pearson')
        print(correlation_matrix)
    else:
        print("数据不足，无法计算相关性")

if __name__ == "__main__":
    prepare_ali_data("salary_data.json", "ali_data.json")